{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "词表大小: 715\n"
     ]
    }
   ],
   "source": [
    "from tokenizers import Tokenizer\n",
    "\n",
    "# 加载 BPE Tokenizer\n",
    "tokenizer = Tokenizer.from_file(\"bpe_tokenizer.json\")\n",
    "\n",
    "# 获取词表大小\n",
    "vocab_size = tokenizer.get_vocab_size()\n",
    "print(\"词表大小:\", vocab_size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "训练数据大小: torch.Size([483, 50]) torch.Size([483, 50])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "def create_training_data(text, tokenizer, seq_length=50):\n",
    "    \"\"\"将文本转换为 Token ID，并生成训练样本\"\"\"\n",
    "    token_ids = tokenizer.encode(text).ids\n",
    "    inputs, targets = [], []\n",
    "\n",
    "    for i in range(len(token_ids) - seq_length):\n",
    "        inputs.append(token_ids[i:i+seq_length])\n",
    "        targets.append(token_ids[i+1:i+seq_length+1])\n",
    "\n",
    "    return torch.tensor(inputs), torch.tensor(targets)\n",
    "\n",
    "# 读取训练数据\n",
    "file_paths = [\"data1.txt\", \"data2.txt\"]\n",
    "text = \"\\n\".join([open(f, \"r\", encoding=\"utf-8\").read() for f in file_paths])\n",
    "\n",
    "# 生成训练样本\n",
    "inputs, targets = create_training_data(text, tokenizer)\n",
    "print(\"训练数据大小:\", inputs.shape, targets.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "module 'torch' has no attribute '_utils'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[15], line 37\u001b[0m\n\u001b[0;32m     34\u001b[0m model \u001b[38;5;241m=\u001b[39m TransformerDecoder(vocab_size, d_model, num_heads, num_layers, dim_feedforward, max_len)\n\u001b[0;32m     36\u001b[0m \u001b[38;5;66;03m# 训练参数\u001b[39;00m\n\u001b[1;32m---> 37\u001b[0m optimizer \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptim\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mAdam\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlr\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.001\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[0;32m     38\u001b[0m scheduler \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39moptim\u001b[38;5;241m.\u001b[39mlr_scheduler\u001b[38;5;241m.\u001b[39mStepLR(optimizer, step_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m, gamma\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0.95\u001b[39m)\n\u001b[0;32m     40\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mtrain\u001b[39m(model, dataloader, epochs\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m20\u001b[39m):\n",
      "File \u001b[1;32mc:\\Users\\16673\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\torch\\optim\\adam.py:45\u001b[0m, in \u001b[0;36mAdam.__init__\u001b[1;34m(self, params, lr, betas, eps, weight_decay, amsgrad, foreach, maximize, capturable, differentiable, fused)\u001b[0m\n\u001b[0;32m     39\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInvalid weight_decay value: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mweight_decay\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     41\u001b[0m defaults \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mdict\u001b[39m(lr\u001b[38;5;241m=\u001b[39mlr, betas\u001b[38;5;241m=\u001b[39mbetas, eps\u001b[38;5;241m=\u001b[39meps,\n\u001b[0;32m     42\u001b[0m                 weight_decay\u001b[38;5;241m=\u001b[39mweight_decay, amsgrad\u001b[38;5;241m=\u001b[39mamsgrad,\n\u001b[0;32m     43\u001b[0m                 maximize\u001b[38;5;241m=\u001b[39mmaximize, foreach\u001b[38;5;241m=\u001b[39mforeach, capturable\u001b[38;5;241m=\u001b[39mcapturable,\n\u001b[0;32m     44\u001b[0m                 differentiable\u001b[38;5;241m=\u001b[39mdifferentiable, fused\u001b[38;5;241m=\u001b[39mfused)\n\u001b[1;32m---> 45\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdefaults\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     47\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m fused:\n\u001b[0;32m     48\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m differentiable:\n",
      "File \u001b[1;32mc:\\Users\\16673\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\torch\\optim\\optimizer.py:266\u001b[0m, in \u001b[0;36mOptimizer.__init__\u001b[1;34m(self, params, defaults)\u001b[0m\n\u001b[0;32m    263\u001b[0m     param_groups \u001b[38;5;241m=\u001b[39m [{\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mparams\u001b[39m\u001b[38;5;124m'\u001b[39m: param_groups}]\n\u001b[0;32m    265\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m param_group \u001b[38;5;129;01min\u001b[39;00m param_groups:\n\u001b[1;32m--> 266\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madd_param_group\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcast\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparam_group\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    268\u001b[0m \u001b[38;5;66;03m# Allows _cuda_graph_capture_health_check to rig a poor man's TORCH_WARN_ONCE in python,\u001b[39;00m\n\u001b[0;32m    269\u001b[0m \u001b[38;5;66;03m# which I don't think exists\u001b[39;00m\n\u001b[0;32m    270\u001b[0m \u001b[38;5;66;03m# https://github.com/pytorch/pytorch/issues/72948\u001b[39;00m\n\u001b[0;32m    271\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_warned_capturable_if_run_uncaptured \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\16673\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\torch\\_compile.py:22\u001b[0m, in \u001b[0;36m_disable_dynamo.<locals>.inner\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m     20\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(fn)\n\u001b[0;32m     21\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minner\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m---> 22\u001b[0m     \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_dynamo\u001b[39;00m\n\u001b[0;32m     24\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_dynamo\u001b[38;5;241m.\u001b[39mdisable(fn, recursive)(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mc:\\Users\\16673\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\torch\\_dynamo\\__init__.py:2\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\n\u001b[1;32m----> 2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m allowed_functions, convert_frame, eval_frame, resume_execution\n\u001b[0;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackends\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mregistry\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m list_backends, register_backend\n\u001b[0;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mconvert_frame\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m replay\n",
      "File \u001b[1;32mc:\\Users\\16673\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\torch\\_dynamo\\allowed_functions.py:26\u001b[0m\n\u001b[0;32m     23\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_functorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdeprecated\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mdeprecated_func\u001b[39;00m\n\u001b[0;32m     24\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_symbolic_trace\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m is_fx_tracing\n\u001b[1;32m---> 26\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m config\n\u001b[0;32m     27\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mexternal_utils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m is_compiling\n\u001b[0;32m     28\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m is_safe_constant, NP_SUPPORTED_MODULES\n",
      "File \u001b[1;32mc:\\Users\\16673\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\torch\\_dynamo\\config.py:51\u001b[0m\n\u001b[0;32m     41\u001b[0m specialize_int \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[0;32m     43\u001b[0m \u001b[38;5;66;03m# Assume these functions return constants\u001b[39;00m\n\u001b[0;32m     44\u001b[0m constant_functions \u001b[38;5;241m=\u001b[39m {\n\u001b[0;32m     45\u001b[0m     torch\u001b[38;5;241m.\u001b[39mjit\u001b[38;5;241m.\u001b[39mis_scripting: \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[0;32m     46\u001b[0m     torch\u001b[38;5;241m.\u001b[39mjit\u001b[38;5;241m.\u001b[39mis_tracing: \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[0;32m     47\u001b[0m     torch\u001b[38;5;241m.\u001b[39m_C\u001b[38;5;241m.\u001b[39m_get_tracing_state: \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[0;32m     48\u001b[0m     torch\u001b[38;5;241m.\u001b[39mfx\u001b[38;5;241m.\u001b[39m_symbolic_trace\u001b[38;5;241m.\u001b[39mis_fx_tracing: \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[0;32m     49\u001b[0m     torch\u001b[38;5;241m.\u001b[39monnx\u001b[38;5;241m.\u001b[39mis_in_onnx_export: \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[0;32m     50\u001b[0m     external_utils\u001b[38;5;241m.\u001b[39mis_compiling: \u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m---> 51\u001b[0m     \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_utils\u001b[49m\u001b[38;5;241m.\u001b[39mis_compiling: \u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[0;32m     52\u001b[0m }\n\u001b[0;32m     54\u001b[0m \u001b[38;5;66;03m# legacy config, does nothing now!\u001b[39;00m\n\u001b[0;32m     55\u001b[0m dynamic_shapes \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\16673\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\torch\\__init__.py:1833\u001b[0m, in \u001b[0;36m__getattr__\u001b[1;34m(name)\u001b[0m\n\u001b[0;32m   1830\u001b[0m     \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mimportlib\u001b[39;00m\n\u001b[0;32m   1831\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m importlib\u001b[38;5;241m.\u001b[39mimport_module(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m.\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;18m__name__\u001b[39m)\n\u001b[1;32m-> 1833\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodule \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m has no attribute \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[1;31mAttributeError\u001b[0m: module 'torch' has no attribute '_utils'"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# 创建 PyTorch 数据加载器\n",
    "dataset = TensorDataset(inputs, targets)\n",
    "dataloader = DataLoader(dataset, batch_size=16, shuffle=True)\n",
    "\n",
    "# 定义 Transformer Decoder\n",
    "class TransformerDecoder(nn.Module):\n",
    "    def __init__(self, vocab_size, d_model, num_heads, num_layers, dim_feedforward, max_len):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, d_model)\n",
    "        self.positional_encoding = nn.Parameter(torch.randn(1, max_len, d_model))\n",
    "        \n",
    "        decoder_layer = nn.TransformerDecoderLayer(d_model, num_heads, dim_feedforward)\n",
    "        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers)\n",
    "        self.fc_out = nn.Linear(d_model, vocab_size)\n",
    "\n",
    "    def forward(self, input_seq, memory):\n",
    "        seq_len = input_seq.size(1)\n",
    "        embedded = self.embedding(input_seq) + self.positional_encoding[:, :seq_len, :]\n",
    "        memory = self.embedding(memory) + self.positional_encoding[:, :memory.size(1), :]\n",
    "        output = self.transformer_decoder(embedded, memory)\n",
    "        return self.fc_out(output)\n",
    "\n",
    "# 超参数\n",
    "d_model = 64  \n",
    "num_heads = 4  \n",
    "num_layers = 3  \n",
    "dim_feedforward = 128  \n",
    "max_len = 100  \n",
    "\n",
    "model = TransformerDecoder(vocab_size, d_model, num_heads, num_layers, dim_feedforward, max_len)\n",
    "\n",
    "# 训练参数\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.95)\n",
    "\n",
    "def train(model, dataloader, epochs=20):\n",
    "    for epoch in range(epochs):\n",
    "        for batch_inputs, batch_targets in dataloader:\n",
    "            optimizer.zero_grad()\n",
    "            output = model(batch_inputs, batch_inputs)\n",
    "            loss = nn.CrossEntropyLoss()(output.view(-1, vocab_size), batch_targets.view(-1))\n",
    "            loss.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 避免梯度爆炸\n",
    "            optimizer.step()\n",
    "        scheduler.step()\n",
    "        print(f\"Epoch {epoch}, Loss: {loss.item()}\")\n",
    "\n",
    "# 训练模型\n",
    "train(model, dataloader, epochs=20)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_text(model, tokenizer, start_text, max_length=50):\n",
    "    \"\"\"用 BPE 词表生成文本\"\"\"\n",
    "    model.eval()\n",
    "    input_tokens = tokenizer.encode(\"<bos> \" + start_text).ids\n",
    "    input_tensor = torch.tensor([input_tokens])\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for _ in range(max_length):\n",
    "            output = model(input_tensor, input_tensor)\n",
    "            next_token = output.argmax(dim=-1)[:, -1].item()\n",
    "            input_tokens.append(next_token)\n",
    "            input_tensor = torch.tensor([input_tokens])\n",
    "\n",
    "            # 如果遇到 <eos> 结束标记，则停止\n",
    "            if tokenizer.decode([next_token]) == \"<eos>\":\n",
    "                break\n",
    "\n",
    "    return tokenizer.decode(input_tokens)\n",
    "\n",
    "# 测试生成\n",
    "print(\"生成文本:\", generate_text(model, tokenizer, \"hello\"))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
